import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

def svd_compress(channel, keep_rate):
    """对单个颜色通道进行压缩，keep_rate为保留的奇异值比例（0-1）"""
    # 分解通道矩阵
    U, s, VT = np.linalg.svd(channel, full_matrices=False)
    # 计算要保留的奇异值数量
    k = int(len(s) * keep_rate)
    k = max(1, k)  # 至少保留1个，避免出错
    # 用前k个奇异值重构通道
    compressed = U[:, :k] @ np.diag(s[:k]) @ VT[:k, :]
    # 确保像素值在0-255范围内
    return np.clip(compressed, 0, 255).astype(np.uint8)

def compress_image(img, keep_rate):
    """对彩色图像压缩，按比例保留每个通道的奇异值"""
    # 分离RGB三个通道
    r = img[:, :, 0]
    g = img[:, :, 1]
    b = img[:, :, 2]

    # 每个通道单独压缩
    r_comp = svd_compress(r, keep_rate)
    g_comp = svd_compress(g, keep_rate)
    b_comp = svd_compress(b, keep_rate)

    # 合并通道得到彩色图像
    return np.stack([r_comp, g_comp, b_comp], axis=2)

def show_images(original, comp_50, comp_25, comp_10, comp_1):
    """展示不同压缩率的效果对比"""
    plt.figure(figsize=(10, 12))  # 设定画布大小

    # 显示原图
    plt.subplot(3, 2, 1)
    plt.imshow(original)
    plt.title("原始图像")
    plt.axis('off')

    # 显示各压缩率图像
    plt.subplot(3, 2, 2)
    plt.imshow(comp_50)
    plt.title("50%压缩率")
    plt.axis('off')

    plt.subplot(3, 2, 3)
    plt.imshow(comp_25)
    plt.title("25%压缩率")
    plt.axis('off')

    plt.subplot(3, 2, 4)
    plt.imshow(comp_10)
    plt.title("10%压缩率")
    plt.axis('off')

    plt.subplot(3, 2, 5)
    plt.imshow(comp_1)
    plt.title("1%压缩率")
    plt.axis('off')

    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    # 加载图像
    face = Image.open('face.png')
    im = np.array(face)  # 转换为NumPy数组

    # 定义要测试的压缩率（对应保留的奇异值比例）
    rates = [0.5, 0.25, 0.1, 0.01]

    # 生成不同压缩率的图像
    comp_50 = compress_image(im, rates[0])
    comp_25 = compress_image(im, rates[1])
    comp_10 = compress_image(im, rates[2])
    comp_1 = compress_image(im, rates[3])

    # 显示对比结果
    show_images(im, comp_50, comp_25, comp_10, comp_1)

    # 保存10%压缩率的图像作为示例
    Image.fromarray(comp_10).save("svd_10pct.jpg")
    print("10%压缩率的图像已保存为：svd_10pct.jpg")